/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package com.example.mcp;

import org.springframework.web.bind.annotation.RestController;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonMappingException;
import com.fasterxml.jackson.databind.ObjectMapper;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;

@RestController()
public class GiteeClientController {

    private static final Logger logger = LoggerFactory.getLogger(GiteeClientController.class);

    @Autowired
    private ChatLanguageModel chatLanguageModel;
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private GiteeClientService giteeApiService;
    @Autowired
    private FederatedQueryEngineService federatedQueryEngineService;
    @Autowired
    private ObjectMapper objectMapper;

    /**
     * 获取用户信息
     * 
     * @return
     */
    @GetMapping("/gitee/user")
    public User getUserInfo() {
        return giteeApiService.getUserInfo();
    }

    @GetMapping("/gitee/sql")
    public List<Map<String, Object>> sqlQuery(@RequestParam(name = "sql") String sql) {
        String answer = federatedQueryEngineService.sqlQuery(sql);
        logger.info("\n" + answer);

        String jsoString = answer.substring(answer.lastIndexOf("```json") + 7);
        jsoString = jsoString.substring(0, jsoString.lastIndexOf("```"));
        logger.info("\n" + jsoString);

        try {
            return objectMapper.readValue(jsoString, new TypeReference<List<Map<String, Object>>>() {});
        } catch (Exception e) {
            throw new RuntimeException(e);   
        }
    }

    /**
     * 测试embedding
     * 
     * @return
     */
    @GetMapping("/gitee/embedding")
    public String embedding() {
        // user in memory embedding store
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        // add some text segments
        TextSegment textSegment = TextSegment.from("我叫放屁大王。我来自中国杭州滨江区。我是一个程序员。我喜欢编程、画画、打游戏。对了,我是男的");
        embeddingStore.add(this.embeddingModel.embed(textSegment).content(), textSegment);
        // create content retriever
        ContentRetriever contentRetriever = new EmbeddingStoreContentRetriever(embeddingStore, embeddingModel);
        // create AI service
        MyAiService myAiService = AiServices.builder(MyAiService.class)
                .chatLanguageModel(chatLanguageModel)
                .contentRetriever(contentRetriever)
                .tools(new MyTools())
                .build();
        // chat
        return myAiService.chat("新问题：我的名字是什么？我住在哪里？");
    }
}
